import argparse
import os
from contextlib import nullcontext
# os.environ["HF_HOME"] = "/media/ubuntu/DC48F14648F120461/hf-home"
import torch
from PIL import Image
from tqdm import tqdm

from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
from spar3d.system import SPAR3D
from spar3d.utils import foreground_crop, get_device


def check_positive(value):
    ivalue = int(value)
    if ivalue <= 0:
        raise argparse.ArgumentTypeError("%s is an invalid positive int value" % value)
    return ivalue


def sanitize_filename(filename):
    """确保文件名合法，去掉特殊字符"""
    return "".join(c if c.isalnum() or c in "._-" else "_" for c in filename)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "image", type=str, nargs="+", help="Path to input image(s) or folder."
    )
    parser.add_argument(
        "--device",
        default=get_device(),
        type=str,
        help=f"Device to use. If no CUDA/MPS-compatible device is found, the baking will fail. Default: '{get_device()}'",
    )
    parser.add_argument(
        "--pretrained-model",
        default="stabilityai/stable-point-aware-3d",
        type=str,
        help="Path to the pretrained model. Could be either a huggingface model id is or a local path. Default: 'stabilityai/stable-point-aware-3d'",
    )
    parser.add_argument(
        "--foreground-ratio",
        default=1.3,
        type=float,
        help="Ratio of the foreground size to the image size. Only used when --no-remove-bg is not specified. Default: 0.85",
    )
    parser.add_argument(
        "--output-dir",
        default="output/",
        type=str,
        help="Output directory to save the results. Default: 'output/'",
    )
    parser.add_argument(
        "--texture-resolution",
        default=1024,
        type=int,
        help="Texture atlas resolution. Default: 1024",
    )
    parser.add_argument(
        "--low-vram-mode",
        action="store_true",
        help=(
            "Use low VRAM mode. SPAR3D consumes 10.5GB of VRAM by default. "
            "This mode will reduce the VRAM consumption to roughly 7GB but in exchange "
            "the model will be slower. Default: False"
        ),
    )

    remesh_choices = ["none"]
    if TRIANGLE_REMESH_AVAILABLE:
        remesh_choices.append("triangle")
    if QUAD_REMESH_AVAILABLE:
        remesh_choices.append("quad")
    parser.add_argument(
        "--remesh_option",
        choices=remesh_choices,
        default="none",
        help="Remeshing option",
    )
    if TRIANGLE_REMESH_AVAILABLE or QUAD_REMESH_AVAILABLE:
        parser.add_argument(
            "--reduction_count_type",
            choices=["keep", "vertex", "faces"],
            default="keep",
            help="Vertex count type",
        )
        parser.add_argument(
            "--target_count",
            type=check_positive,
            help="Selected target count.",
            default=2000,
        )
    parser.add_argument(
        "--batch_size", default=1, type=int, help="Batch size for inference"
    )
    args = parser.parse_args()

    # Ensure args.device contains cuda
    devices = ["cuda", "mps", "cpu"]
    if not any(args.device in device for device in devices):
        raise ValueError("Invalid device. Use cuda, mps or cpu")

    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

    device = args.device

    print("Device used: ", device)

    model = SPAR3D.from_pretrained(
        args.pretrained_model,
        config_name="config.yaml",
        weight_name="model.safetensors",
        low_vram_mode=args.low_vram_mode,
    )
    model.to(device)
    model.eval()

    images = []

    def handle_image(image_path):
        """处理单张图像并存储结果"""
        # **修改部分：获取原始文件名并构造输出路径**
        image_name = os.path.splitext(os.path.basename(image_path))[0]
        sanitized_name = sanitize_filename(image_name)
        image_output_dir = os.path.join(output_dir, sanitized_name)
        os.makedirs(image_output_dir, exist_ok=True)

        # **修改部分：裁剪图像并以原始文件名保存**
        image = Image.open(image_path).convert("RGBA")
        image = foreground_crop(image, args.foreground_ratio)
        cropped_image_path = os.path.join(image_output_dir, f"{sanitized_name}.png")
        image.save(cropped_image_path)

        # 将图像及输出信息存储
        images.append((image, sanitized_name, image_output_dir))

    # 处理输入路径的逻辑
    idx = 0
    for image_path in args.image:
        if os.path.isdir(image_path):
            image_paths = [
                os.path.join(image_path, f)
                for f in os.listdir(image_path)
                if f.endswith((".png", ".jpg", ".jpeg"))
            ]
            for path in image_paths:
                handle_image(path)
                idx += 1
        else:
            handle_image(image_path)
            idx += 1

    # vertex_count 参数计算逻辑
    if hasattr(args, "reduction_count_type"):
        vertex_count = (
            -1
            if args.reduction_count_type == "keep"
            else (
                args.target_count
                if args.reduction_count_type == "vertex"
                else args.target_count // 2
            )
        )
    else:
        vertex_count = -1  # 默认行为

    for i, (image, sanitized_name, image_output_dir) in enumerate(tqdm(images)):
        with torch.no_grad():
            with (
                torch.autocast(device_type=device, dtype=torch.bfloat16)
                if "cuda" in device
                else nullcontext()
            ):
                mesh, glob_dict = model.run_image(
                    [image],
                    bake_resolution=args.texture_resolution,
                    remesh=args.remesh_option,
                    vertex_count=vertex_count,
                    return_points=True,
                )

        # **修改部分：保存文件名与输入一致**
        mesh_path = os.path.join(image_output_dir, f"{sanitized_name}.glb")
        points_path = os.path.join(image_output_dir, f"{sanitized_name}.ply")
        mesh.export(mesh_path, include_normals=True)
        glob_dict["point_clouds"][0].export(points_path)

        # **修改部分：打印输出路径**
        print(f"Saved: {mesh_path}, {points_path}, {os.path.join(image_output_dir, f'{sanitized_name}.png')}")
